package etxt2db.annotators;

import java.net.URI;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import org.apache.uima.UimaContext;
import org.apache.uima.analysis_component.AnalysisComponent;
import org.apache.uima.analysis_component.JCasAnnotator_ImplBase;
import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.examples.SourceDocumentInformation;
import org.apache.uima.jcas.JCas;
import org.apache.uima.jcas.tcas.Annotation;
import org.apache.uima.resource.ResourceInitializationException;

import etxt2db.api.ClassificationExecutor;
import etxt2db.api.ClassificationExecutor.ClassifiedSpan;
import etxt2db.api.ClassificationModel;
import etxt2db.api.Document;
import etxt2db.corpus.DocumentLoader;
import etxt2db.corpus.MyAnnotation;
import etxt2db.mappings.AnnotationToKeyMapper;
import etxt2db.serialization.ClassificationModelSerializer;


/**
 * The UIMA annotator wrapper used by UCompare to classify documents 
 * from the E-txt2db classifier.
 * 
 * @author Gonçalo Simões
 * @author Rui Lageira
 *
 */
public class ComponentAnnotator extends JCasAnnotator_ImplBase {
	
	String modelPath;
	List<String> attributes;
	
	String splitterClassName;
	String tokenizerClassName;
	
	ClassificationModelSerializer serial;
	ClassificationModel model;
	
	DocumentLoader loader;
	
	AnnotationToKeyMapper mapper;

	/**
	 * @see AnalysisComponent#initialize(UimaContext)
	 */
	public void initialize(UimaContext aContext) throws ResourceInitializationException {
		super.initialize(aContext);
		
		// Get the configuration parameters
		modelPath = (String) aContext.getConfigParameterValue("ModelPath");	
		String[] attributesArray = (String[]) aContext.getConfigParameterValue("Attributes");
		attributes = new ArrayList<String>(Arrays.asList(attributesArray));
		String[] mappingsContent = (String[]) aContext.getConfigParameterValue("Mappings");
		splitterClassName = (String) aContext.getConfigParameterValue("Splitter");
		tokenizerClassName = (String) aContext.getConfigParameterValue("Tokenizer");
		
		// Load the model
		serial = new ClassificationModelSerializer();
		model = serial.deserializeClassificationModel(modelPath);
		
		loader = new DocumentLoader(model.getSplitter(), model.getTokenizer());
		
		mapper = new AnnotationToKeyMapper(new ArrayList<String>(Arrays.asList(mappingsContent)));
		
	}
	
	@Override
	public void process(JCas jcas) throws AnalysisEngineProcessException {

		try {
			// Load the etxt2db document. It will gather the informations from the previous context of this UCompare component.
			List<Document> documents;
			documents = loader.load(jcas.getDocumentText());

			Iterator<Document> documentIter = documents.iterator();
			
			// Get the first (and only one) document.
			Document document = null;			
			if(documentIter.hasNext()) {
				document = documentIter.next();
			}
			
			// Get the document URI to insert in etxt2db Document
			Iterator<?> it = jcas.getAnnotationIndex(SourceDocumentInformation.type).iterator();
			
			String docUri = null;
		    if (it.hasNext()) {
		      SourceDocumentInformation srcDocInfo = (SourceDocumentInformation) it.next();
		      docUri = srcDocInfo.getUri();
		      document.setUri(new URI(docUri));
		    }
		    
			// Iterate and gather annotations from a preceding workflow.
			Iterator<?> annotationIter = jcas.getAnnotationIndex().iterator();
			while (annotationIter.hasNext()) {
				
				Annotation annotation = (Annotation) annotationIter.next();
				
				// The list of sub-types correspondent to this generic UIMA annotation.
				List<String> keys = mapper.getKey(annotation);
				
				for(String type : keys) {
					
					MyAnnotation myAnn = new MyAnnotation(annotation, type, mapper.getAnnotations(annotation, type));
					
					document.addAnnotation(myAnn);
				}
			}

			// CLASSIFY
			
			ClassificationExecutor exec = new ClassificationExecutor();
			
			Map<String, List<ClassifiedSpan>> spans = exec.getClassifiedSpans(documents, model, attributes);
			
			for(Entry<String, List<ClassifiedSpan>> spansEntry : spans.entrySet()) {
				List<ClassifiedSpan> sentenceSpans = spansEntry.getValue();
				for(ClassifiedSpan span : sentenceSpans) {
					NamedEntityTyped annotation = new NamedEntityTyped(jcas);
					annotation.setBegin(span.getStart());
					annotation.setEnd(span.getEnd());
					annotation.setTag(span.getType());
					annotation.addToIndexes();
				}
			}
			
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}

}